# ===============================================================
# file_name: Trajectory_data_generator2.py
# function: Maneuvering target Trajectory Data Generator
# For common Airport Surveillance Radar
# The detection distance coverage: 926m~37040m
# For common civil aircraft (including helicopter)
# The velocity coverage: -340m/s~340m/s
# The turn rate coverage: -10°/s~10°/s changing interval 0.1°/s
# ===============================================================

import numpy as np
import random as rd
from numpy.linalg import cholesky
from scipy.io import savemat
# Sampling interval
T = 0.1
t_segment = 5
# The velocity coverage of starting point:
velocity_max = 340
velocity_min = -340
# The distance coverage of starting point:
distance_max = 37040
distance_min = 926
# The turn rage coverage:
turn_rate_max = 10
turn_rate_min = -10
turn_rate_interval = 0.1

# trajectory data generator
def trajectory_create(data_len, transition_matrix):
    # Transition noise(Acceleration noise m/s^2)
    state_n = rd.uniform(8.0, 13.0)  # .uniform(a,b)函数用于生成一个指定范围内的均匀分布的随机浮点数
    state_var = np.square(state_n)  # 方差
    T2 = np.power(T, 2)
    T3 = np.power(T, 3)
    T4 = np.power(T, 4)
    var_m = np.array([[T4 / 4, 0, 0, 0], [0, T4 / 4, 0, 0], [0, 0, T2, 0], [0, 0, 0, T2]]) * state_var
    chol_var = cholesky(var_m)

    # 产生一个data_len*4的矩阵，每个元素都是0.
    data = np.array([[0 for i in range(4)] for j in range(data_len)], 'float64')
    distance_change_max = velocity_max * T * data_len
    # Starting point
    sp_distance = rd.uniform(distance_min + distance_change_max, distance_max - distance_change_max)
    # rd.random()函数用于生成一个[0,1)范围内的均匀分布的随机浮点数
    # 该行代码表示生成一个(-180°, 180°)范围内的均匀分布的角度
    sp_direction = (rd.random() - 0.5) * 2 * np.pi
    distance_x = sp_distance * np.cos(sp_direction)  # Target X direction position
    distance_y = sp_distance * np.sin(sp_direction)  # Target Y direction position

    sp_velocity = rd.uniform(velocity_min, velocity_max)
    velocity_direction = (rd.random()-0.5) * 2 * np.pi
    velocity_x = sp_velocity * np.cos(velocity_direction)   # Target X direction velocity
    velocity_y = sp_velocity * np.sin(velocity_direction)   # Target Y direction velocity

    x_true_store = np.array([[distance_x, distance_y, velocity_x, velocity_y]], 'float64')
    for i in range(data_len):
        data[i, :] = x_true_store
        x_true_store = np.dot(x_true_store, transition_matrix)
    data_n = data + np.dot(np.random.randn(data_len, 4), chol_var)   # 加噪声
    return data, data_n

def trajectory_batch_generator(batch_size, data_len):
    # Initialize
    trajectory_result = np.zeros((batch_size, data_len, 4), dtype='float64')
    observation_result = np.zeros((batch_size, data_len, 2), dtype='float64')
    transition_matrix = np.zeros((batch_size, 4, 4), dtype='float64')
    for i in range(batch_size):
        turn_rate = rd.randint(int(turn_rate_min/turn_rate_interval), int(turn_rate_max/turn_rate_interval))*turn_rate_interval
        if turn_rate == 0:
            F_matrix = np.array([
                            [1, 0, T, 0],
                            [0, 1, 0, T],
                            [0, 0, 1, 0],
                            [0, 0, 0, 1]
                            ])
        else:
            w = turn_rate * np.pi/180   # 转为弧度
            F_matrix = np.array([
                            [1, 0, np.sin(w*T)/w, (np.cos(w*T)-1)/w],
                            [0, 1, (1-np.cos(w*T))/w, np.sin(w*T)/w],
                            [0, 0, np.cos(w*T), -np.sin(w*T)],
                            [0, 0, np.sin(w*T), np.cos(w*T)]
                            ])
        F_matrix = np.transpose(F_matrix, [1, 0])   # 将矩阵进行转置
        dt, dtn = trajectory_create(data_len, F_matrix)
        trajectory_result[i] = dt

        # Observation
        # Observational noise
        distance_noise = rd.uniform(8.0, 13.0)  # distance noise
        distance_var = np.square(distance_noise)     # distance noise variance
        azimuth_noise = rd.uniform(7.0, 9.0)    # azimuth noise
        azimuth_var = np.square(azimuth_noise/1000)  # azimuth noise variance

        observation_result[i, :, 0] = np.arctan2(dtn[:, 1], dtn[:, 0]) + np.random.normal(0, azimuth_var, data_len)
        observation_result[i, :, 1] = np.sqrt(np.square(dtn[:, 0])+np.square(dtn[:, 1])) + np.random.normal(0, distance_var, data_len)
        transition_matrix[i, :, :] = F_matrix
    return trajectory_result, observation_result, transition_matrix

def trajectory_batch_generator2(batch_size, data_len):
    # Initialize
    # Observation noise in X Y direction
    trajectory_result = np.array([[[0 for i in range(4)] for j in range(data_len)] for k in range(batch_size)],
                      'float64')  # Initialization of trajectory
    observation_result = np.array([[[0 for i in range(2)] for j in range(data_len)] for k in range(batch_size)],
                     'float64')  # Initialization of observation
    transition_matrix = np.array([[[0 for i in range(4)] for j in range(4)] for k in range(batch_size)], 'float64')

    for i in range(batch_size):

        turn_rate = rd.randint(int(turn_rate_min / turn_rate_interval), int(turn_rate_max / turn_rate_interval)) * turn_rate_interval
        if turn_rate == 0:
            F_matrix = np.array([[1, 0, T, 0], [0, 1, 0, T], [0, 0, 1, 0], [0, 0, 0, 1]])
        else:
            w = turn_rate * np.pi / 180  # turn rate
            F_matrix = np.array([[1, 0, np.sin(w * T) / w, (np.cos(w * T) - 1) / w],
                            [0, 1, -(np.cos(w * T) - 1) / w, np.sin(w * T) / w],
                            [0, 0, np.cos(w * T), -np.sin(w * T)],
                            [0, 0, np.sin(w * T), np.cos(w * T)]])

        F_matrix = np.transpose(F_matrix, [1, 0])
        dt, dtn = trajectory_create(data_len, F_matrix)
        trajectory_result[i] = dt

        # Observations
        # observation noise in X Y direction
        X_noise = rd.uniform(10.0, 50.0)  # X direction noise
        X_var = np.square(X_noise)
        dtn[:, 0] = dtn[:, 0] + np.random.normal(0, X_var, data_len)

        Y_noise = rd.uniform(10.0, 50.0)  # X direction noise
        Y_var = np.square(Y_noise)
        dtn[:, 1] = dtn[:, 1] + np.random.normal(0, Y_var, data_len)

        observation_result[i, :, 0] = np.arctan2(dtn[:, 1], dtn[:, 0])  # Azimuth
        observation_result[i, :, 1] = np.sqrt(np.square(dtn[:, 0]) + np.square(dtn[:, 1]))  # Distance
        transition_matrix[i, :, :] = F_matrix
    return trajectory_result, observation_result, transition_matrix

if __name__ == '__main__':
    trajectory_result, observation_result, transition_matrix = trajectory_batch_generator(100, 50)
    # 保存成mat数据
    file_name = './matlab/my_data.mat'
    savemat(file_name, {'my_trajectory': trajectory_result, 'my_observation': observation_result, 'my_transition_matrix': transition_matrix})
